package org.example.es;

import com.google.common.collect.ImmutableMap;
import io.searchbox.client.JestClient;
import io.searchbox.client.JestClientFactory;
import io.searchbox.client.JestResultHandler;
import io.searchbox.client.config.HttpClientConfig;
import io.searchbox.core.*;
import io.searchbox.indices.CreateIndex;
import io.searchbox.indices.mapping.PutMapping;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.conn.ssl.NoopHostnameVerifier;
import org.apache.http.conn.ssl.SSLConnectionSocketFactory;
import org.apache.http.nio.conn.SchemeIOSessionStrategy;
import org.apache.http.nio.conn.ssl.SSLIOSessionStrategy;
import org.apache.http.ssl.SSLContextBuilder;
import org.apache.http.ssl.TrustStrategy;
import org.elasticsearch.common.StopWatch;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.mapper.Mapping;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import javax.net.ssl.HostnameVerifier;
import javax.net.ssl.SSLContext;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class JestTest {

    private JestClient client;

    private static final AtomicInteger counter = new AtomicInteger(0);

    private static final int batchSize = 10000;

    private static final int threadCount = 10;

    private Map<Long, StopWatch> map = new ConcurrentHashMap<>();

    private JestClientFactory factory;

    @Before
    public void before() throws Exception {

        SSLContext sslContext = new SSLContextBuilder().loadTrustMaterial(null, new TrustStrategy() {
            public boolean isTrusted(X509Certificate[] arg0, String arg1) throws CertificateException {
                return true;
            }
        }).build();

        // skip hostname checks
        HostnameVerifier hostnameVerifier = NoopHostnameVerifier.INSTANCE;

        SSLConnectionSocketFactory sslSocketFactory = new SSLConnectionSocketFactory(sslContext, hostnameVerifier);
        SchemeIOSessionStrategy httpsIOSessionStrategy = new SSLIOSessionStrategy(sslContext, hostnameVerifier);

        factory = new JestClientFactory();
        factory.setHttpClientConfig(new HttpClientConfig
                .Builder("https://192.168.206.122:9200")
                .connTimeout(50000)
                .readTimeout(50000)
                .defaultSchemeForDiscoveredNodes("https")
                .defaultCredentials("admin", "admin")
                .sslSocketFactory(sslSocketFactory)
                .httpsIOSessionStrategy(httpsIOSessionStrategy)
                .multiThreaded(true)
                //Per default this implementation will create no more than 2 concurrent connections per given route
                // .defaultMaxTotalConnectionPerRoute(<YOUR_DESIRED_LEVEL_OF_CONCURRENCY_PER_ROUTE>)
                // and no more 20 connections in total
                // .maxTotalConnection(<YOUR_DESIRED_LEVEL_OF_CONCURRENCY_TOTAL>)
                .build());
    }

    @Test
    public void deleteIndex() {

    }

    @Test
    public void createIndex() throws Exception {
        // 创建索引
        Settings settings = Settings.builder()
//                .put("index.number_of_shards", "5")
//                .put("index.number_of_replicas", "0")
//                .put("index.refresh_interval", "120s")
//                .put("index.translog.durability", "async")
//                .put("index.translog.sync_interval", "120s")
//                .put("index.translog.sync_interval", "30m")
//                .put("index.translog.flush_threshold_size", "1g")
//                .put("index.merge.scheduler.max_thread_count", "1")
                .build();

        CreateIndex.Builder builder = new CreateIndex.Builder("ep_tag_001");
        builder.settings(settings.toString());

        String mappings = "{\n" +
                "  \"mappings\": {\n" +
                "    \"properties\": {\n" +
                "      \"address\": {\n" +
                "        \"type\": \"text\",\n" +
                "        \"fields\": {\n" +
                "          \"raw\": {\n" +
                "            \"type\": \"text\",\n" +
                "            \"analyzer\": \"ik_smart\"\n" +
                "          }\n" +
                "        }\n" +
                "      },\n" +
                "      \"belong_industry\": {\n" +
                "        \"type\": \"text\",\n" +
                "        \"fields\": {\n" +
                "          \"raw\": {\n" +
                "            \"type\": \"text\",\n" +
                "            \"analyzer\": \"ik_smart\"\n" +
                "          }\n" +
                "        }\n" +
                "      },\n" +
                "      \"company_name\": {\n" +
                "        \"type\": \"text\",\n" +
                "        \"fields\": {\n" +
                "          \"raw\": {\n" +
                "            \"type\": \"text\",\n" +
                "            \"analyzer\": \"ik_smart\"\n" +
                "          }\n" +
                "        }\n" +
                "      }\n" +
                "    }\n" +
                "  }\n" +
                "}";
        builder.mappings(mappings);

        // 字符串类型的字段默认会分词
        // Map<String, Object> maps = new HashMap<>();

        factory.getObject().execute(builder.build());
    }

    @Test
    public void putIndex() throws Exception{
        PutMapping putMapping = new PutMapping.Builder(
                "ep_tag_001",
                "ep_tag_001",
                "{ \"my_type\" : { \"properties\" : { \"message\" : {\"type\" : \"string\", \"store\" : \"yes\"} } } }"
        ).build();
        factory.getObject().execute(putMapping);
    }

    @Test
    public void insertData() throws Exception {
        Bulk.Builder builder = new Bulk.Builder()
                .defaultIndex("ep_tag_001");

//        Map<String, Object> map = new HashMap<>();
//        map.put("company_name", "国泰新点");
//        map.put("address", "张家港市杨舍镇沙洲湖科创园");
//        map.put("belong_industry", "软件行业");
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("company_name", "国泰新点")
                .put("address", "张家港市杨舍镇沙洲湖科创园")
                .put("belong_industry", "软件行业")
            .build()).build());


        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("company_name", "国泰华盛")
                .put("address", "张家港市杨舍镇人民中路")
                .put("belong_industry", "制造业")
                .build()).build());

        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("company_name", "沙钢")
                .put("address", "张家港市锦丰镇")
                .put("belong_industry", "制造业")
                .build()).build());

        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("company_name", "国泰航空")
                .put("address", "香港金钟道88号太古广场一座33楼")
                .put("belong_industry", "服务业")
                .build()).build());

        factory.getObject().execute(builder.build());

        Thread.sleep(Integer.MAX_VALUE);
    }

    @Test
    public void insertData_tag() throws Exception {
        Bulk.Builder builder = new Bulk.Builder()
                .defaultIndex("ep_tag_002");

//        Map<String, Object> map = new HashMap<>();
//        map.put("company_name", "国泰新点");
//        map.put("address", "张家港市杨舍镇沙洲湖科创园");
//        map.put("belong_industry", "软件行业");
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "企业名称")
                .put("tag_value", "国泰新点")
                .build()).build());
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "企业名称")
                .put("tag_value", "国泰华盛")
                .build()).build());
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "企业名称")
                .put("tag_value", "沙钢")
                .build()).build());
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "企业名称")
                .put("tag_value", "国泰航空")
                .build()).build());

        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "地址")
                .put("tag_value", "张家港市杨舍镇沙洲湖科创园")
                .build()).build());
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "地址")
                .put("tag_value", "张家港市杨舍镇人民中路")
                .build()).build());
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "地址")
                .put("tag_value", "张家港市锦丰镇")
                .build()).build());
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "地址")
                .put("tag_value", "香港金钟道88号太古广场一座33楼")
                .build()).build());

        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "所属行业")
                .put("tag_value", "软件行业")
                .build()).build());
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "所属行业")
                .put("tag_value", "制造业")
                .build()).build());
        builder.addAction(new Index.Builder(new ImmutableMap.Builder<>()
                .put("tag_name", "所属行业")
                .put("tag_value", "服务业")
                .build()).build());

        factory.getObject().execute(builder.build());

        Thread.sleep(Integer.MAX_VALUE);
    }

    @Test
    public void searchIndex() throws Exception{
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder
                .query(QueryBuilders.matchPhraseQuery("address", "张家港"))
        ;

        Search search = new Search.Builder(searchSourceBuilder.toString())
                // multiple index or types can be added.
                .addIndex("ep_tag_001")
                .build();

        SearchResult result = factory.getObject().execute(search);
        log.info(result.getJsonString());
    }

    @Test
    public void searchTemplateIndex() throws Exception{
        String query = "{\n" +
                "    \"params\": {\n" +
                "        \"query_string\" : \"张家港 国泰\"" +
                "    }\n" +
                "}";

        Search search = new Search.TemplateBuilder(query)
                // multiple index or types can be added.
                .addIndex("ep_tag_001")
                .build();

        SearchResult result = factory.getObject().execute(search);
        log.info(result.getJsonString());
    }

    @Test
    public void bulkIndex() throws Exception {

        ExecutorService executorService = Executors.newFixedThreadPool(threadCount);
        CountDownLatch latch = new CountDownLatch(threadCount);

        StopWatch stopWatch = new StopWatch();
        stopWatch.start("totalTimes");

//        // 创建索引
//        Settings settings = Settings.builder()
//                .put("number_of_shards", "1")
//                .put("number_of_replicas", "0")
//                .put("refresh_interval", "-1")
//                .put("translog.durability", "async")
//                .put("translog.sync_interval", "120s")
//                .put("translog.sync_interval", "30m")
//                .put("translog.flush_threshold_size", "1g")
//                .build();
//        factory.getObject().execute(new CreateIndex.Builder("ep_tag_001").settings(settings.toString()).build());

        for (int i = 0; i < threadCount; i++) {
            executorService.submit(new Runnable() {
                @Override
                public void run() {
                    try {
                        StopWatch stopWatch = new StopWatch();
                        stopWatch.start(Thread.currentThread().getName());

                        Bulk.Builder builder = new Bulk.Builder()
                                .defaultIndex("ep_tag_001");

                        // 循环10次 10 * 10000
                        for (int j = 0; j < 10; j++) {
                            for (int i = 0; i < batchSize; i++) {
                                Map<String, Object> map = new HashMap<>();
                                int tagId = counter.getAndIncrement();
                                map.put("tag_id", tagId);
                                map.put("tag_name", "name" + tagId);
                                map.put("tag_value", "value" + tagId);
                                builder.addAction(new Index.Builder(map).build());
                            }
//                        factory.getObject().executeAsync(builder.build(), new JestResultHandler<BulkResult>() {
//                            @Override
//                            public void completed(BulkResult bulkResult) {
//                                stopWatch.stop();
//                                log.info("result = {}, {}条, time = {}", bulkResult.isSucceeded(), batchSize, stopWatch.shortSummary());
//                            }
//
//                            @Override
//                            public void failed(Exception e) {
//                                e.printStackTrace();
//                            }
//                        });

                            factory.getObject().execute(builder.build());
                        }
//                        for (int i = 0; i < batchSize; i++) {
//                            Map<String, Object> map = new HashMap<>();
//                            int tagId = counter.getAndIncrement();
//                            map.put("tag_id", tagId);
//                            map.put("tag_name", "name" + tagId);
//                            map.put("tag_value", "value" + tagId);
//                            builder.addAction(new Index.Builder(map).build());
//                        }
////                        factory.getObject().executeAsync(builder.build(), new JestResultHandler<BulkResult>() {
////                            @Override
////                            public void completed(BulkResult bulkResult) {
////                                stopWatch.stop();
////                                log.info("result = {}, {}条, time = {}", bulkResult.isSucceeded(), batchSize, stopWatch.shortSummary());
////                            }
////
////                            @Override
////                            public void failed(Exception e) {
////                                e.printStackTrace();
////                            }
////                        });
//
//                        BulkResult bulkResult = factory.getObject().execute(builder.build());
//                        stopWatch.stop();
//                        log.info("result = {}, {}条, time = {}", bulkResult.isSucceeded(), batchSize, stopWatch.shortSummary());
                        latch.countDown();
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
            });
        }

        latch.await();
        stopWatch.stop();
        log.info("{}条，耗时：{}", threadCount * batchSize, stopWatch.shortSummary());

        Thread.sleep(Integer.MAX_VALUE);
    }

    @After
    public void close() throws Exception {
        // client.close();
    }
}
